package icu.funkye.confine.limiter.product.impl;

import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import icu.funkye.confine.common.loader.LoadLevel;
import icu.funkye.confine.limiter.product.AbstractLimiter;
import icu.funkye.confine.limiter.product.Limiter;
import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;

/**
 * @author jianbin.chen
 */
@LoadLevel(name = "bucket")
public class BucketLimiter extends AbstractLimiter {

    Bucket bucket;

    @Override
    public Limiter create(String resource, int tps) {
        return create(resource, tps, 1);
    }

    @Override
    public Limiter create(String resource, int tps, long seconds) {
        currentSeconds = seconds;
        Bandwidth limit = Bandwidth.simple(tps, Duration.ofSeconds(seconds));
        bucket = Bucket4j.builder().addLimit(limit).build();
        this.resource = resource;
        return this;
    }

    @Override
    public Limiter setTps(int tps) {
        Bandwidth limit = Bandwidth.simple(tps, Duration.ofSeconds(currentSeconds));
        bucket = Bucket4j.builder().addLimit(limit).build();
        currentSeconds = tps;
        return this;
    }

    @Override
    public int getTps() {
        return currentTps;
    }

    @Override
    public boolean tryAcquire(int permits, long timeout, TimeUnit unit) {
        boolean tryAcquire = false;
        if (timeout > 0) {
            CompletableFuture<Boolean> future = CompletableFuture.supplyAsync(() -> {
                boolean result;
                while (true) {
                    result = bucket.tryConsume(permits);
                    if (!result) {
                        try {
                            Thread.sleep(100);
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    } else {
                        break;
                    }
                }
                return result;
            });
            try {
                tryAcquire = future.get(timeout, unit);
            } catch (Exception e) {
            }
        } else {
            tryAcquire = bucket.tryConsume(permits);
        }
        return tryAcquire;
    }

    @Override
    public boolean tryAcquire(int permits) {
        return tryAcquire(permits, 0, TimeUnit.MILLISECONDS);
    }

    @Override
    public boolean tryAcquire() {
        return tryAcquire(1);
    }

}
